package com.cnpc.bds.sdk.sink.kafka.executor;

import com.cnpc.bds.sdk.sink.kafka.meta.ConsumerCfg;
import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;

import java.time.Duration;
import java.time.temporal.ChronoUnit;
import java.util.LinkedList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * sink-kafka上游消费任务——kafka-consumer.
 *
 * @author ssp
 * @since 1.0
 */
@Slf4j
public abstract class UpstreamConsumer extends Thread {

    private final ConsumerCfg consumerCfg;
    private Consumer<String, String> consumer;

    private final String topic;

    public UpstreamConsumer(String name, ConsumerCfg consumerCfg) {
        super(name);
        this.consumerCfg = consumerCfg;
        this.topic = consumerCfg.getTopicMeta().getTopic();
    }

    public UpstreamConsumer(ConsumerCfg consumerCfg) {
        this.consumerCfg = consumerCfg;
        this.topic = consumerCfg.getTopicMeta().getTopic();
    }

    @Override
    public void run() {
        exec();
    }

    protected void exec() {
        //TODO poll：每分区每线程并行拉取数据 ---> 转消费者线程组实现.

        init0();
        do {
            try {
                List<Object> msgList = poll();
                consume(msgList);
            } catch (Exception e) {
                log.warn("Consume exception!", e);
            }
        } while (true);
    }

    private void init0() {
        consumer = new KafkaConsumer<>(consumerCfg.properties());

        final String topic = consumerCfg.getTopicMeta().getTopic();
        final List<PartitionInfo> partitionInfos = consumer.partitionsFor(topic);
        final List<TopicPartition> partitions = partitionInfos.stream()
                .map(partitionInfo -> new TopicPartition(partitionInfo.topic(), partitionInfo.partition()))
                .collect(Collectors.toList());

        consumer.assign(partitions);
    }

    protected abstract void consume(List<Object> msgList);

    private List<Object> poll() {

        final List<Object> result = new LinkedList<>();
        final ConsumerRecords<String, String> records = consumer.poll(Duration.of(6, ChronoUnit.SECONDS));
        records.iterator()
                .forEachRemaining(r -> result.add(r.value()));
        log.debug("<--- poll topic={} count={}: {}", this.topic, result.size(), result);
        return result;
    }

    /**
     * 下游任务链挂钩.
     *
     * @return hook
     */
    public abstract Object downStreamHook();

}
